package core

import (
	"encoding/binary"
	"golang.org/x/net/bpf"
	"golang.org/x/net/ipv4"
	"net"
)

func SetRawConnBPF(rc *ipv4.RawConn, filter []bpf.Instruction) error {
	assembled, err := bpf.Assemble(filter)
	if err != nil {
		return err
	}
	return rc.SetBPF(assembled)
}

func CreateBPFDropAll() []bpf.Instruction {
	return []bpf.Instruction{
		bpf.RetConstant{Val: 0x0},
	}
}

func SetBPFDropAll(rc *ipv4.RawConn) error {
	filter := CreateBPFDropAll()
	assembled, err := bpf.Assemble(filter)
	if err != nil {
		return err
	}
	return rc.SetBPF(assembled)
}

func CreateBPFICMP(id uint16, dstIP net.IP) []bpf.Instruction {
	ipUint32 := binary.BigEndian.Uint32(dstIP.To4())
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(28), SkipTrue: 9}, // ip header + icmp or ip header + icmp + ip header + icmp

		bpf.LoadMemShift{Off: 0},                                                    // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                                           // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(0), SkipTrue: 1},                // echo reply
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(11), SkipTrue: 2, SkipFalse: 5}, // time exceed

		bpf.LoadIndirect{Off: 4, Size: 2},                                           // echo reply identifier
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(id), SkipTrue: 2, SkipFalse: 3}, // echo reply

		bpf.LoadIndirect{Off: 24, Size: 4},                           // time exceed dst ip address
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: ipUint32, SkipFalse: 1}, // dst ip

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDPICMP(dstIP net.IP) []bpf.Instruction {
	ipUint32 := binary.BigEndian.Uint32(dstIP.To4())
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(56), SkipTrue: 9}, // ip header + icmp + ip header + udp

		bpf.LoadMemShift{Off: 0},                                       // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                              // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(3), SkipTrue: 1},   // destination unreachable
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(11), SkipFalse: 3}, // time exceed

		bpf.LoadIndirect{Off: 17, Size: 1},                             // ip protocol
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(17), SkipFalse: 1}, // udp

		bpf.LoadIndirect{Off: 24, Size: 4},                           // dst ip address
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: ipUint32, SkipFalse: 1}, // echo reply

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFTCPICMP(dstIP net.IP) []bpf.Instruction {
	ipUint32 := binary.BigEndian.Uint32(dstIP.To4())
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(56), SkipTrue: 8}, // ip header + icmp + ip header + tcp

		bpf.LoadMemShift{Off: 0},                                       // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                              // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(11), SkipFalse: 5}, // time exceed

		bpf.LoadIndirect{Off: 17, Size: 1},                            // ip protocol
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(6), SkipFalse: 1}, // tcp

		bpf.LoadIndirect{Off: 24, Size: 4},                           // dst ip address
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: ipUint32, SkipFalse: 1}, // echo reply

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDP(pid uint16) []bpf.Instruction {
	return []bpf.Instruction{

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 4, Size: 2},                                // udp length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(20), SkipTrue: 7}, // 20 bytes

		bpf.LoadIndirect{Off: 8, Size: 4},                              // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x6e747261, SkipFalse: 5}, // ntra

		bpf.LoadIndirect{Off: 12, Size: 4},                             // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x63652121, SkipFalse: 3}, // ce!!

		bpf.LoadIndirect{Off: 16, Size: 2},                              // pid
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(pid), SkipFalse: 1}, // 7

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFTCP(dstIP net.IP, srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	ipUint32 := binary.BigEndian.Uint32(dstIP.To4())
	tcpFlagAck := 0x10
	tcpFlagRst := 0x04
	tcpFlagSynAck := 0x12
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(40), SkipTrue: 14}, // ip header + tcp header

		bpf.LoadIndirect{Off: 12, Size: 4},                            // src ip address
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: ipUint32, SkipFalse: 12}, // target ip

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                               // tcp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 9}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 8},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                               // tcp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 6}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 5},    // max

		bpf.LoadIndirect{Off: 13, Size: 1},                                          // tcp flag
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(tcpFlagAck), SkipTrue: 2},       // ACK包
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagRst), SkipTrue: 1},     // RST包,
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagSynAck), SkipFalse: 1}, // SYN+ACK包

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}
